from vocab import WordVocab  # 假设 vocab.py 文件名为 vocab.py

vocab_path = "vocab.txt"  # 你的 vocab.txt 文件路径
text_vocab_output_path = "vocab_text_representation.txt" # 新的文本词汇表输出路径

try:
    vocab = WordVocab.load_vocab(vocab_path)  # 加载 vocab.txt 文件
    print(f"成功加载词汇表，词汇表大小: {len(vocab)}")

    print("\n词汇表前 20 个词语:")
    for i in range(50):
        print(f"{i + 1}: {vocab.itos[i]}")  # 打印前 20 个词语

    # ---  新增代码：遍历词汇表并保存到文本文件 ---
    print(f"\n正在将词汇表保存到文本文件: {text_vocab_output_path}")
    with open(text_vocab_output_path, 'w', encoding='utf-8') as f:
        for token in vocab.itos:  # 遍历词汇表中的词语列表 (itos)
            f.write(token + '\n')  # 将每个词语写入文本文件，每行一个词语
    print(f"词汇表文本文件已保存到: {text_vocab_output_path}")

except FileNotFoundError:
    print(f"错误: 词汇表文件 {vocab_path} 未找到，请先运行 vocab.py 生成词汇表。")
except Exception as e:
    print(f"加载词汇表时发生错误: {e}")
    import traceback

    traceback.print_exc()
